package utils

import (
	"errors"
	"fmt"
	"net"
	"sync/atomic"
	"time"
)

// Proxy links two network connections, allowing the two remote ends
// to communicate. The function blocks until the proxied connections
// are closed or an error occurs. The returned LogEntry provides a
// summary of the exchange.
func Proxy(first, second net.Conn, timeout time.Duration) *LogEntry {
	record := &LogEntry{StartTime: time.Now().UTC()}
	record.First.RemoteAddr = first.RemoteAddr()
	record.First.LocalAddr = first.LocalAddr()
	record.Second.RemoteAddr = second.RemoteAddr()
	record.Second.LocalAddr = second.LocalAddr()

	err := pump(
		first,
		second,
		timeout,
		&record.First.SentBytes,
		&record.Second.SentBytes)

	// Wait for the pumps to die and provide an error.
	record.ShutdownCause = <-err
	record.EndTime = time.Now().UTC()
	return record
}

// LogEntry provides information about a proxied session between two
// net.Conns.
type LogEntry struct {
	// Times are UTC.
	StartTime, EndTime time.Time

	// Identity and activity of either side of the session.
	First, Second struct {
		RemoteAddr, LocalAddr net.Addr

		// Bytes sent by this end of the session.
		SentBytes uint64
	}

	// The error which caused the session to be torn down.
	ShutdownCause error
}

func (l *LogEntry) String() string {
	return fmt.Sprintf(
		`Started: %s
Ended: %s
Link: %s -> %s -> %s -> %s
Sent: %d bytes
Received: %d bytes
Shutdown cause: %s`,
		l.StartTime, l.EndTime,
		l.First.RemoteAddr, l.First.LocalAddr,
		l.Second.LocalAddr, l.Second.RemoteAddr,
		l.First.SentBytes, l.Second.SentBytes, l.ShutdownCause)
}

type pumpst struct {
	a, b    net.Conn
	buf     []byte
	sent    *uint64
	w       *watchdog
	timeout time.Duration
	err     chan error
}

const inactivityTimeout = 10e9

func pump(a, b net.Conn, timeout time.Duration, aSent, bSent *uint64) chan error {
	w := newWatchdog(timeout)
	err := make(chan error, 2)
	p1 := &pumpst{a, b, make([]byte, 10*1024), aSent, w, timeout, err}
	p2 := &pumpst{b, a, make([]byte, 10*1024), bSent, w, timeout, err}
	go p1.run()
	go p2.run()
	return err
}

func (p *pumpst) run() {
	for {
		n, err := p.read()
		if err != nil {
			p.err <- err
			p.a.Close()
			p.b.Close()
			return
		}

		err = p.write(n)
		if err != nil {
			p.err <- err
			p.a.Close()
			p.b.Close()
			return
		}
	}
}

func (p *pumpst) read() (int, error) {
	for {
		p.a.SetReadDeadline(time.Now().Add(p.timeout))
		n, err := p.a.Read(p.buf)
		if n == 0 {
			if err2, ok := err.(net.Error); ok && err2.Timeout() {
				// Check watchdog, we may have to continue
				if p.w.expired() {
					return 0, errors.New("No client activity")
				}
				continue
			} else {
				return 0, err
			}
		} else {
			// Network activity, so reset watchdog.
			p.w.reset()
			return n, nil
		}
	}
	panic("unreachable")
}

func (p *pumpst) write(n int) error {
	wb := p.buf[:n]
	for {
		p.b.SetWriteDeadline(time.Now().Add(p.timeout))
		nw, err := p.b.Write(wb)
		if err != nil {
			if err, ok := err.(net.Error); ok && err.Timeout() {
				if p.w.expired() {
					return errors.New("No client activity")
				}
				wb = wb[nw:]
				continue
			}
			return err
		} else {
			p.w.reset()
			atomic.AddUint64(p.sent, uint64(n))
			return nil
		}
	}
	panic("unreachable")
}
